from PIL import Image
import numpy as np 
from sklearn.cluster import KMeans
import matplotlib.pyplot as plt 

def show_img_compar(img_1, img_2):
    f, ax = plt.subplots(1, 2, figsize=(10,10))
    ax[0].imshow(img_1)
    ax[1].imshow(img_2)
    ax[0].axis('off') #hide the axis
    ax[1].axis('off')
    f.tight_layout()
    plt.show()

def palette(clusters):
    width=300
    palette = np.zeros((50, width, 3), np.uint8)
    steps = width/clusters.cluster_centers_.shape[0]
    for idx, centers in enumerate(clusters.cluster_centers_): 
        palette[:, int(idx*steps):(int((idx+1)*steps)), :] = centers
    return palette

image = Image.open("./1.jpeg")
image = np.array(image)

print(image.shape)

image_flatten = image.reshape((-1, 3))

print(image_flatten.shape)

unique, count = np.unique(image_flatten, axis=0, return_counts=True)

print(unique)
print(count)

clt = KMeans(n_clusters=5)
# clt.fit(image_flatten)
# print(clt.labels_)
# print(clt.cluster_centers_)

clt_1 = clt.fit(image_flatten)
show_img_compar(image, palette(clt_1))
 
